from utils import *
from joblib import load
from pyod.models.ocsvm import OCSVM
import os
import sys
import numpy as np
from threading import Lock
import queue
from models.base_model import BaseModel, ModelConfig
from utils import sys_logger
from config import config_singleton

sys.path.append(
    os.path.abspath(os.path.join(os.path.dirname("__file__"), '..')))


class IMU_AnomalyDetector(BaseModel):
    """
    IMU 陀螺异常检测模型, 使用OCSVM算法, 输入为IMU 1s增量数据, 输出为异常检测结果.
    IMU 加表异常检测模型，使用OCSVM算法, 输入为IMU 1s增量数据, 输出为异常检测结果.
    """

    def __init__(self, config: ModelConfig):
        super().__init__(config)

        self.input_dim = 0
        self.input_update_flag = dict()
        for k, v in config.inputs.items():
            # (flag, offset)
            self.input_update_flag[k] = [False, self.input_dim]
            self.input_dim += len(v)

        self.output_dim = len(config.outputs["var_names"])
        
        self.bus_gaps = config.input_gaps

        # dim0 为batchsize
        self.input_data = np.zeros((1, self.input_dim)) # 初始化为零

        self.output_data_queue = queue.Queue(maxsize=config.input_output_queue_size)
        self.input_data_queue = queue.Queue(maxsize=config.input_output_queue_size)

        # 线程锁
        self.lock = Lock()
        self.thread_stop = False

        self.run_cnt = 0
        self.train_mode = config.online_training_mode           # True为在线训练模式，False为离线推理模式

        self.models = []
        if self.train_mode:
            self.models.append(OCSVM())   # Gyro
            self.models.append(OCSVM())   # Accel
        else:
            self.load_weights(config.model_weigths_file)

        self.train_data = np.zeros((self._config.online_training_interval, self.input_dim))
        self.output_table_num = self._config.outputs["table_num"]


    def update_input_data(self, msgbean: MsgBean):
        """
        更新输入数据.
        根据数据所在的位置和所在的表号，更新输入数据.
        """
        table_num = msgbean.table_num

        # 如果数据busgap不符合要求，则将该消息丢弃, 相差0.004s以上
        if abs(msgbean.bus_gap - self.bus_gaps[table_num]) > 0.004:
            return

        # 获取当前表的偏移量, 存储在inputdata中的偏移量
        offset = self.input_update_flag[table_num][1]

        # cnt = len(self._config.inputs[table_num])
        with self.lock:
            # 根据数组中的值作为字典的key, 取出对应key的value
            # numpy 可以使用切片的方式。
            # self.input_data[0][offset: offset +
            #                    cnt] = msgbean.vars[self._config.inputs[table_num]]
            for index in self._config.inputs[table_num]:
                self.input_data[0][offset] = msgbean.vars[index]
                offset += 1
            # 更新标志位
            self.input_update_flag[table_num][0] = True

        if all(flag for [flag, _] in self.input_update_flag.values()):
            # TODO 考虑是否使用带时间，推理结束后无法对应于那个时间的数据
            self.input_data_queue.put([msgbean.vars[1], self.input_data.copy()])
            # 数据放入了队列中，便可以把标志清理掉。
            # 先将标志重置，再进行推理
            with self.lock:
                for k, v in self.input_update_flag.items():
                    v[0] = False


    def run(self):
        """
        模型推理线程.
        一帧数据进来，就进行模型推理，并将推理结果放入输出队列中.
        需要启动多线程，一个线程负责接收数据。
        """
        while not self.thread_stop:
            # 阻塞式等待数据的到来。
            t, input = self.input_data_queue.get()

            if input is None or t is None:
                continue

            if np.all(input == 0):
                continue

            # train_mode == True 则是收集数据，待训练
            if self.train_mode:
                self.train_data[self.run_cnt] = input
                self.run_cnt += 1
                # 预留接口，可以通过设置train_mode = True，控制程序再次进入训练模式
                if self.run_cnt % self._config.online_training_interval == 0:
                    self.models[0].fit(self.train_data[...,0:5])
                    self.models[1].fit(self.train_data[...,5:])
                    self.train_mode = False       # 进入预测模式
            else:
                vals = np.zeros(self.output_dim, np.float32)
                # 输出数据数组前部分存储输入值，后部分存储异常检测值
                vals[0: self.input_dim] = input[0, :]
                
                vals[self.input_dim] = self.models[0].predict(input[:, 0:5])
                vals[self.input_dim + 1] = self.models[1].predict(input[:, 5:])

                # 将时间拼接起来。
                self.output_data_queue.put([t, vals])
                self.run_cnt += 1
        
        sys_logger.info(f"{self._config.name} Model inference thread stopped. Total predict data num {self.run_cnt}.")


    def get_output_data(self):
        """
        获取输出数据,
        用于后续数据重组和网络发送。
        """
        if self.output_data_queue.empty():
            vals = [None, None]
        else:
            vals = self.output_data_queue.get_nowait()
        
        return vals


    def stop(self):
        """
        停止模型推理线程.
        通过设置标志位为True, 线程会退出循环, 停止运行.
        """
        self.thread_stop = True
        self.input_data_queue.put([None, None])
        self.output_data_queue.put([None, None])


    def train(self, X_train, y_train):
        """
        训练模型.
        暂时未实现，后续根据需要可以在线训练的模型.
        """
        for model in self.models:
            model.fit(X_train, y_train)


    def load_weights(self, weights_path):
        """
        加载模型参数.
        不同模型加载参数的方式不一样，需要根据具体模型进行修改.
        """
        if len(weights_path) != 2:
            raise ValueError("weights_path should be a list of two elements.")

        for tmp_path in weights_path:

            if config_singleton.Debug:
                real_path = os.path.join(os.path.dirname(__file__), '..', 'model_files',tmp_path)
            else:
                real_path = os.path.join(os.path.dirname(os.path.realpath(sys.executable)), 'model_files', tmp_path)

            if not os.path.exists(real_path):
                raise FileNotFoundError("IMU Anomaly Detector Model weights file not found." + real_path)
            else:
                self.models.append(load(real_path))
           

    # 同时获取输出变量的名字和数值
    def get_output_varnames(self):
        
        return self._config.outputs["var_names"]
        
    
    def get_model_name(self):

        return self._config.name




if __name__ == '__main__':

    config = ModelConfig()
    model = IMU_AnomalyDetector(config)

    print("")
